import socket
import threading
import time
import json
from typing import Dict, Any, Optional, Callable
from ..server.protocol import Protocol, MessageType

class DatabaseClientWithCallback:
    """支持回调的数据库客户端"""
    
    def __init__(self, host: str = "localhost", port: int = 8888):
        self.host = host
        self.port = port
        self.socket = None
        self.connected = False
        self.running = False
        self.session_id = None
        self.user_info = None
        
        # 回调函数
        self.on_query_result: Optional[Callable[[Dict[str, Any]], None]] = None
        self.on_error: Optional[Callable[[str, str], None]] = None
        self.on_success: Optional[Callable[[str], None]] = None
        
    def set_callbacks(self, 
                     on_query_result: Optional[Callable[[Dict[str, Any]], None]] = None,
                     on_error: Optional[Callable[[str, str], None]] = None,
                     on_success: Optional[Callable[[str], None]] = None):
        """设置回调函数"""
        self.on_query_result = on_query_result
        self.on_error = on_error
        self.on_success = on_success
        
    def connect(self) -> bool:
        """连接到服务器"""
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.connect((self.host, self.port))
            self.connected = True
            self.running = True
            
            # 启动接收线程
            receive_thread = threading.Thread(target=self._receive_messages)
            receive_thread.daemon = True
            receive_thread.start()
            
            return True
        except Exception as e:
            if self.on_error:
                self.on_error("CONNECTION_ERROR", f"连接失败: {e}")
            return False
            
    def disconnect(self):
        """断开连接"""
        self.running = False
        self.connected = False
        if self.socket:
            self.socket.close()
            
    def login(self, username: str, password: str) -> bool:
        """用户登录"""
        if not self.connected:
            if self.on_error:
                self.on_error("NOT_CONNECTED", "请先连接到服务器")
            return False
            
        try:
            message = Protocol.create_login_request(username, password)
            self.socket.send(message)
            time.sleep(0.5)  # 等待响应
            return self.session_id is not None
        except Exception as e:
            if self.on_error:
                self.on_error("LOGIN_ERROR", f"登录失败: {e}")
            return False
            
    def logout(self) -> bool:
        """用户登出"""
        if not self.session_id:
            return True
            
        try:
            message = Protocol.encode_message(MessageType.LOGOUT, {}, self.session_id)
            self.socket.send(message)
            self.session_id = None
            self.user_info = None
            return True
        except Exception as e:
            if self.on_error:
                self.on_error("LOGOUT_ERROR", f"登出失败: {e}")
            return False
            
    def execute_sql(self, sql: str) -> bool:
        """执行SQL语句"""
        if not self.session_id:
            if self.on_error:
                self.on_error("NOT_LOGGED_IN", "请先登录")
            return False
            
        try:
            message = Protocol.encode_message(MessageType.SQL_QUERY, {"sql": sql}, self.session_id)
            self.socket.send(message)
            return True
        except Exception as e:
            if self.on_error:
                self.on_error("SQL_ERROR", f"执行SQL失败: {e}")
            return False
            
    def _receive_messages(self):
        """接收服务器消息的线程"""
        buffer = b""
        
        while self.running and self.connected:
            try:
                data = self.socket.recv(4096)
                if not data:
                    break
                    
                buffer += data
                
                # 处理完整的消息
                while len(buffer) >= 4:
                    msg_length = int.from_bytes(buffer[:4], byteorder='big')
                    if len(buffer) < 4 + msg_length:
                        break
                        
                    message_data = buffer[:4 + msg_length]
                    buffer = buffer[4 + msg_length:]
                    
                    message = Protocol.decode_message(message_data)
                    if message:
                        self._handle_server_message(message)
                        
            except Exception as e:
                if self.running and self.on_error:
                    self.on_error("RECEIVE_ERROR", f"接收消息时出错: {e}")
                break
                
    def _handle_server_message(self, message: Dict[str, Any]):
        """处理服务器消息"""
        msg_type = message.get("type")
        data = message.get("data", {})
        session_id = message.get("session_id")
        
        if msg_type == MessageType.RESPONSE.value:
            if data.get("success"):
                # 处理成功响应
                if "user_info" in data:  # 登录响应
                    self.session_id = session_id
                    self.user_info = data["user_info"]
                    if self.on_success:
                        self.on_success(f"登录成功！欢迎 {self.user_info['username']}")
                elif "data" in data and "columns" in data["data"]:  # 查询结果
                    if self.on_query_result:
                        self.on_query_result(data["data"])
                else:  # 其他成功响应
                    if self.on_success:
                        self.on_success(data.get('message', '操作成功'))
            else:
                if self.on_error:
                    self.on_error("OPERATION_FAILED", data.get('message', '操作失败'))
                    
        elif msg_type == MessageType.ERROR.value:
            if self.on_error:
                error_code = data.get('error_code', 'UNKNOWN')
                error_message = data.get('error_message', '未知错误')
                self.on_error(error_code, error_message)